from Crypto.Cipher import AES
from Crypto.Hash import SHA256, SHA1
from Crypto.Random import get_random_bytes
import hashlib
import time
import pickle

'''
    《原版协议分步骤运行》
    本代码将所有阶段按照函数进行分开
    1. 异或需要用int类型的值进行，且异或后的结果值的类型为int
    2. sha函数生成的值均为字节类型，若要参与计算，需要将该值设置为int，可以使用  int.from_bytes(x.digest(), byteorder='big')来转换
       sha函数的输入需要为byte，所以需要将其他类型换算为byte，在此项目中，需要将异或后的值（int）转换为byte进行计算，
       可以使用 hwA.to_bytes((hwA.bit_length() + 7) // 8, byteorder='big')  进行转换
       其他数值，byte转int可以用 int.from_bytes(message, byteorder='big')  进行转换
       ----1 byte = 8 bit
    3. get_random_bytes(1)，是随机生成1个字节的值，类型为byte
    4. 可写函数进行int和byte转换，不至于代码混乱
'''

# 全局变量
g = 3


# idA = b"User_A_8150"
# idB = b"User_B_6623"
# pw = b"ssx2016360."


# int to byte

def intToByte(message):
    result = message.to_bytes((message.bit_length() + 7) // 8, byteorder='big')
    return result


# byte to int
def byteToInt(message):
    result = int.from_bytes(message, byteorder='big')
    return result


# 哈希函数
def H(data):
    # 检查输入的数据类型是否为整数
    if isinstance(data, int):
        # 如果是整数，将其转换为字节数组
        # (data.bit_length() + 7) // 8 计算字节数组的长度
        # byteorder='big' 指定字节顺序为大端序
        data = data.to_bytes((data.bit_length() + 7) // 8, byteorder='big')
    # 检查输入的数据类型是否为字符串
    elif isinstance(data, str):
        # 如果是字符串，将其编码为字节数组
        data = data.encode()

    # 使用 SHA256 算法对字节数组进行哈希
    hash_obj = SHA1.new(data)
    # 将哈希值的十六进制表示转换为整数并返回
    return int(hash_obj.hexdigest(), 16)


def H1(message):
    return H(message)  # 将sha1生成的字节值转换为int


def H2(data):
    # 检查输入的数据类型是否为整数
    if isinstance(data, int):
        data = data
    # 检查输入的数据类型是否为字符串
    elif isinstance(data, str):
        # 如果是字符串，将其编码为字节数组
        data_byte = data.encode()
        data = byteToInt(data_byte)  # 字节转int
    elif isinstance(data, bytes):
        data = byteToInt(data)  # 字节转int

    hash_obj = data ^ (data >> 16) ^ (data >> 8)
    # 将哈希值的十六进制表示转换为整数并返回
    return hash_obj


def H3(message):
    return H(message)  # 将sha1生成的字节值转换为int


# 加密函数
def encrypt(key, message):
    plaintext = pickle.dumps(message)  # 将message的元组修改为字节类型
    cipher = AES.new(key, AES.MODE_GCM)  # 创建一个加密对象
    nonce = cipher.nonce  # 创建一次性数字，用于保证在密钥和明文相同的情况下，加密的结果不一致
    ciphertext, tag = cipher.encrypt_and_digest(plaintext)  # ciphertext是密文，tag是认证标签，保证密文的完整性和真实性
    return ciphertext, nonce, tag


# 解密函数
def decrypt(key, ciphertext, nonce, tag):
    cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
    t = cipher.decrypt(ciphertext)  # 解密，结果为字节
    plaintext = pickle.loads(t)  # 将解密后的明文修转化为元组
    try:
        cipher.verify(tag)  # 验证tag是否相同，相同则未被修改
        return plaintext[0], plaintext[1], plaintext[2]  # 将元组中的数据拿出，在此代码中为加密后的数据XB、YB、TB等值
    except ValueError:
        print("key is not valid")
        return False


# 时间戳验证
def timestamp_is_fresh(timestamp):
    now = time.time()
    if now - timestamp > 60:  # 60s内皆为新鲜
        return False
    else:
        return True


# 注册阶段
def generate_password_file(password, idd):
    s = get_random_bytes(1)  # 随机生成s
    hw = H1(password) ^ byteToInt(s)  # 将随机生成的随机数s从字节转换成int,异或只能int类型数据
    D = H2(s) ^ byteToInt(idd)  # idd是字节，需要转换成int类型进行异或
    return hw, D  # 其中hw、D均为int类型


# 阶段1：A设备进行参数生成
def new_session_UA(password_file, T):
    rA = byteToInt(get_random_bytes(1))  # 生成rA，并转换为int类型
    hw, D = password_file
    XA = H2(hw) ^ H2(rA)
    YA = g ** rA
    return XA, YA, T, rA


def new_session_UB(password_file, id, T, receive_XA, receive_YA, receive_TA):
    rB = byteToInt(get_random_bytes(1))
    hw, D = password_file
    XB = H2(hw) ^ H2(rB)
    YB = g ** rB
    if not timestamp_is_fresh(receive_TA):
        print("TA has expired")
    else:
        tkB = receive_XA ^ D ^ byteToInt(id) ^ H2(rB) ^ H2(H3(receive_YA ** rB))
        # 共20位，差4位 用0补齐
        tkB_byte = intToByte(tkB)

        # 补齐字节
        padding = (8 - len(tkB_byte) % 8) % 8  # 计算补齐的字节数，8的倍数
        tkB_byte += b'\x00' * padding  # 补齐后的值

        # tkB = b't2Fb9ZvN8hsPQcRxg6oWnm45EDjKLYHa'
        print(f"tkB={tkB_byte}")
        encty_file_CB = encrypt(tkB_byte, (T, hw, rB))  # 其中包含CB、nonce、tag
    return XB, YB, T, encty_file_CB, rB, tkB_byte


def validate_and_compute_key_UA(receive_XB, receive_YB, receive_TB, receive_encty_file_CB, XA, YA, rA, TA, id,
                                password_file):
    hw, D = password_file
    if not timestamp_is_fresh(receive_TB):
        print("TB has expired")
    else:
        tkA = receive_XB ^ D ^ byteToInt(id) ^ H2(rA) ^ H2(H3(receive_YB ** rA))
        tkA_byte = intToByte(tkA)

        # 补齐字节
        padding = (8 - len(tkA_byte) % 8) % 8
        tkA_byte += b'\x00' * padding

        # tkA = b't2Fb9ZvN8hsPQcRxg6oWnm45EDjKLYHa'
        print(f"tkA={tkA_byte}")
        CB, nonce, tag = receive_encty_file_CB
        dec_TB, dec_hwB, dec_rB = decrypt(tkA_byte, CB, nonce, tag)
        if dec_TB != receive_TB:
            print("dec_TB does not match receive_TB, is not valid")
        else:
            if receive_YB != g ** dec_rB or receive_XB != H2(dec_hwB) ^ H2(dec_rB):
                print("YB does not match or XB does not match")
            else:
                encry_file_CA = encrypt(tkA_byte, (TA, hw, rA))
                sidA = str(XA) + str(YA) + str(receive_XB) + str(receive_YB)  # XA、YA、XB、YB都是经过哈希或者幂运算的，都是int类型
                # print("sidA=", sidA)
                # print("A.tkb=", str(tkA_byte).encode('utf-8'))
                # print("A.sidb=", sidA.encode('utf-8'))
                KA = H(str(tkA_byte).encode('utf-8') + sidA.encode('utf-8'))
                print('KA=', KA)
                return encry_file_CA, KA


def validate_and_compute_key_UB(receive_XA, receive_YA, receive_TA, receive_encry_file_CA, XB, YB, tkB):  # 此处传的tkB实际是tkB_byte
    CA, nonce, tag = receive_encry_file_CA
    dec_TA, dec_hwa, dec_ra = decrypt(tkB, CA, nonce, tag)
    if dec_TA != receive_TA:
        print("dec_TA does not match receive_TA, is not valid")
    elif receive_YA != g ** dec_ra or receive_XA != H2(dec_hwa) ^ H2(dec_ra):
        print("YA does not match or XA does not match")
    else:
        sidB = str(receive_XA) + str(receive_YA) + str(XB) + str(YB)
        # print("sidB=", sidB)
        # print("B.tkb=",str(tkB).encode('utf-8'))
        # print("B.sidb=", sidB.encode('utf-8'))
        KB = H(str(tkB).encode('utf-8') + sidB.encode('utf-8'))
        print("KB=", KB)
        return KB


def session():
    # 设备uA注册
    pwA = b"ssx2016360."
    idA = b"User_A_8150"
    password_file_A = generate_password_file(pwA, idA)

    # 设备UB注册
    pwB = b"ssx2016360."
    idB = b"User_B_6623"
    password_file_B = generate_password_file(pwB, idB)

    # 第一阶段 new_session_UA
    TA = time.time()
    XA, YA, Ta, rA = new_session_UA(password_file_A, TA)

    # 第二阶段 new_session_UB
    TB = time.time()
    XB, YB, Tb, encty_file_CB, rB, tkB = new_session_UB(password_file_B, idB, TB, XA, YA, Ta)

    # 第三阶段 UA验证生成key
    CA, KA = validate_and_compute_key_UA(XB, YB, Tb, encty_file_CB, XA, YA, rA, TA, idA, password_file_A)

    # 第四阶段 UB验证生成key
    KB = validate_and_compute_key_UB(XA, YA, Ta, CA, XB, YB, tkB)

    assert KA == KB, "Session keys do not match!"
    print("Key exchange successful!")


if __name__ == '__main__':
    # 计算单次运行时间
    start_time = time.time()
    session()
    end_time = time.time()
    print("运行时间：", end_time - start_time)

    # 计算运行10000次平均的消耗的时间
    # i = 0
    # sum_time = 0
    # while (i < 1000):
    #     start_time = time.time()
    #     session()
    #     end_time = time.time()
    #     run_time = end_time - start_time
    #     sum_time += run_time
    #     i = i + 1
    #     print("第" + str(i) + "运行时间：", run_time)
    # print("运行1000次的时间：", sum_time)
    # print("运行1000次的平均时间：", sum_time / 1000)
